import json
import sys
from typing import List, Optional, Dict, Any, Tuple, Union
from uuid import UUID

from langchain.memory import ConversationTokenBufferMemory
from langchain.tools.render import render_text_description
from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.language_models import BaseChatModel
from langchain_core.output_parsers import PydanticOutputParser, StrOutputParser
from langchain_core.outputs import GenerationChunk, ChatGenerationChunk, LLMResult
from langchain_core.prompts import PromptTemplate
from langchain_core.tools import StructuredTool
import urllib.parse
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field, ValidationError
import webbrowser
from langchain.agents import initialize_agent, AgentType
from langchain_community.chat_models import ChatOllama

# 自定义工具：打开浏览器搜索
def open_google_search(query: str) -> str:
    encoded_query = urllib.parse.quote(query)
    url = f"https://www.google.com/search?q={encoded_query}"
    webbrowser.open(url)
    return f"已在浏览器中打开搜索：{query}"

# 输入模型定义
class SearchInput(BaseModel):
    query: str = Field(description="需要搜索的关键词")

# 封装为 LangChain 工具
search_tool = StructuredTool.from_function(
    func=open_google_search,
    name="google_search",
    description="在Google中搜索关键词",
    args_schema=SearchInput,
    return_direct=True
)

# 使用 Ollama 本地模型
llm = ChatOllama(model="qwen3:8b", temperature=0)

# 初始化 Agent
agent = initialize_agent(
    tools=[search_tool],
    llm=llm,
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,  # Ollama 也支持函数调用 Agent 类型
    verbose=True
)

# 测试运行
response = agent.run("帮我搜索 github")
print(response)